#!/usr/bin/python
# -*- coding: utf-8 -*-

from __future__ import with_statement

from pylab import meshgrid, linspace
from pylab import xlabel, ylabel, title
from numpy import asarray, seterr

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.axes3d import Axes3D


class DerivativePlots(object):

    """ use to visualize results from pricing engine. """

    __slots__ = [
        'm_axis_leyend',
        'm_curve_leyends',
        'm_colors',
        'm_exit_time',
        'm_fig',
        'm_xcoor',
        'm_yvals',
        'm_vol',
        'm_vol_analy',
        'm_ptitle',
        'm_posi',
        'm_line_style',
        'm_yhist',
        'm_time_elapse',
        ]

    def __init__(self):
        """ initilialize all the variable neede for plotting. """

        (self.m_xcoor, self.m_yvals) = (None, None)
        (self.m_vol, self.m_vol_analy) = (None, None)

        (self.m_time_elapse, self.m_exit_time) = (None, -1.0)
        (self.m_ptitle, self.m_line_style, self.m_curve_leyends) = \
            (None, None, None)

        self.m_yhist = None
        self.m_posi = 1
        self.m_axis_leyend = {}
        self.m_colors = ['blue', 'green', 'red', 'yellow']

        # Set error handling to known state.

        seterr(all='ignore')
        self.m_fig = plt.figure(figsize=(11.5, 5.0))

    def set_curves(self, xcoor, ycoors):
        """ ordinate and possible several absise. """

        (self.m_xcoor, self.m_yvals) = (xcoor, ycoors)

    def set_history(self, xcoor, yhist, time_elapse):
        """ use to plot time stepper history. """

        (self.m_xcoor, self.m_yhist, self.m_time_elapse) = (xcoor, yhist,
                time_elapse)

    def set_vol_surface(self, xcoor, y_coor, vol, vol_analy=None):
        """ 3D plot. """

        (self.m_xcoor, self.m_yvals, self.m_vol, self.m_vol_analy) = (xcoor,
                                                      y_coor, vol, vol_analy)

    def set_curve_leyends(self, names):
        """ give some meaningfull description of the curve. """

        self.m_curve_leyends = names

    def set_leyend_posi(self, posi):
        """ left, right, bot, top. """

        self.m_posi = posi

    def set_axis_leyend(self, axis_leyend):
        """ describe the plotting. """

        self.m_axis_leyend = axis_leyend

    def set_title(self, ptitle):
        """ what is it you are plotting. """

        self.m_ptitle = ptitle

    def set_colors(self, colors):
        """ look at matplot for valid colors. """

        self.m_colors = colors

    def set_line_style(self, line_style):
        """ look at matplot for valid colors. """

        self.m_line_style = line_style

    def set_exit_time(self, exp_time):
        """ not implemented yet. """

        self.m_exit_time = float(exp_time)

    def plot(self, axis=None):
        """ plot 2 or 3 D. """

        if axis is None:
            axis = self.m_fig.gca()

        handle = []
        for (i, y_coor) in enumerate(self.m_yvals):
            (hand, ) = axis.plot(self.m_xcoor, y_coor, '-',
                                 color=self.m_colors[i])
            handle.append(hand)

        if 0 < len(self.m_curve_leyends):
            axis.legend(handle, self.m_curve_leyends, self.m_posi)

        self._add_axis_leyends()

        plt.show()

    def plot_history(self):
        """ use this function to plot pde history. """

        (x_grid, t_grid) = meshgrid(self.m_xcoor, linspace(0, \
                                    self.m_time_elapse, self.m_yhist.shape[0]))
        assert x_grid.shape == self.m_yhist.shape

        rect = self.m_fig.add_subplot(1, 2, 1).get_position()
        ax1 = Axes3D(self.m_fig, rect)
        ax1.plot_surface(x_grid, t_grid, asarray(self.m_yhist, 'float16'))
        #self._add_axis_leyends()

        axis2 = self.m_fig.add_subplot(1, 2, 2)
        self.plot(axis2)

    def plot_vol_surfaces(self):
        """ self describing name. """

        assert self.m_xcoor is not None
        assert self.m_yvals is not None
        assert self.m_vol is not None

        # ---- First subplot

        if self.m_vol_analy is not None:
            rect = self.m_fig.add_subplot(1, 2, 1).get_position()
            axis = Axes3D(self.m_fig, rect)
        else:
            axis = Axes3D(self.m_fig)

        axis.plot_surface(self.m_xcoor, self.m_yvals, asarray(self.m_vol,
                        'float16'), rstride=1, cstride=1,
                        linewidth=0, antialiased=False)


        axis.set_xlabel('Stoke Price')
        axis.set_ylabel('Time to Expiration')
        axis.set_zlabel('Calibrated Volatility')

        # ---- Second subplot

        if self.m_vol_analy is not None:
            rect = self.m_fig.add_subplot(1, 2, 2).get_position()
            axis2 = Axes3D(self.m_fig, rect)

            axis2.plot_surface(self.m_xcoor, self.m_yvals, asarray(
                self.m_vol_analy, 'float16'), rstride=1, cstride=1,
                linewidth=0, antialiased=False)

            axis2.set_xlabel('Stoke Price')
            axis2.set_ylabel('Time to Expiration')
            axis2.set_zlabel('Analytical Volatility')

        plt.show()

    def plot_greeks(self):
        """ greeks. """

        ax1 = self.m_fig.add_subplot(1, 2, 1)
        ax1.plot(self.m_xcoor, self.m_yvals[0])
        self._add_axis_leyends()

        axis2 = self.m_fig.add_subplot(1, 2, 2)
        axis2.plot(self.m_xcoor, self.m_yvals[1])
        axis2.plot(self.m_xcoor, self.m_yvals[2])
        self.m_axis_leyend['y'] = 'Delta / Gamma'
        self._add_axis_leyends()

        self.exit_if_requested()
        plt.show()

    def exit_if_requested(self):
        """ this method is not working yet. """

        if self.m_exit_time > 0:
            pass

        # look at : http://matplotlib.sourceforge.net/examples/
        #                 animation/animate_decay_tk_blit.py
        #" it is not working "
        #time.wait(self.m_exit_time)
        #sys.exit()

    def _add_axis_leyends(self):
        """ is there any thing to put on the leyend. """

        if self.m_axis_leyend.has_key('x'):
            xlabel(self.m_axis_leyend['x'])

        if self.m_axis_leyend.has_key('y'):
            ylabel(self.m_axis_leyend['y'])

        if self.m_axis_leyend.has_key('z'):
            ylabel(self.m_axis_leyend['z'])

        if self.m_ptitle is not None:
            title(self.m_ptitle)


if __name__ == '__main__':
    from numpy import ones, exp, linspace

    deri_plot = DerivativePlots()

    (plot_curves, plot_surface) = (False, True)

    if plot_curves:
        x_coor = asarray(linspace(0, 1, 100))
        y2 = 2 * ones(x_coor.shape)
        y1 = exp(x_coor)

        deri_plot.set_curves(x_coor, [y1, y2])
        deri_plot.set_curve_leyends(['one', 'two'])
        deri_plot.set_axis_leyend({'x': 'price'})
        deri_plot.set_leyend_posi(2)
        deri_plot.plot()

    if plot_surface:
        x_coor = asarray(linspace(-2, 2, 50))
        (X, Y) = meshgrid(x_coor, x_coor)
        Z = exp(X ** 2 + Y ** 2)

        deri_plot.set_vol_surface(X, Y, Z, Z)
        deri_plot.plot_vol_surfaces()

